Repository navigation
refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location - #368
Merged
Merged
Conversation
… point location The JAX-path pure_callback previously ran the full scipy_delaunay pipeline on the host every likelihood evaluation (qhull triangulation + find_simplex on the data and split grids + numpy assembly) - 40-320 ms of host-serial work that scales with over-sampling and serializes per vmap lane. The callback now returns only the qhull triangulation and its adjacency arrays (tri.neighbors plus one incident simplex per vertex); point location, dual areas and split points run inside the JIT program, on GPU and batched across vmap lanes in production. Point location uses the same visibility-walk algorithm find_simplex itself uses (nearest-vertex start, step through the neighbor opposite the most negative barycentric weight, hull exit -> nearest-vertex fallback), as a masked fixed-trip lax.fori_loop chunked over query points. Mappings are exact: 100.000% identity vs find_simplex on production HST grids; full-pipeline JIT log-evidence matches eager to 5.4e-9 (previous JAX path: 5.2e-9), pinned value 29110.92085793 holds. The numpy/scipy path is byte-identical and the jax_delaunay / jax_delaunay_matern return contracts (shapes, dtypes, -1 padding) are unchanged, so no downstream edits are required. PyAutoArray#367 (phase 1 of 2; phase 2 adds the autolens_workspace_test parity suite). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 9, 2026
Collaborator
Author
|
Phase 2 parity-guard PR (merge after this one): PyAutoLabs/autolens_workspace_test#155 |
This was referenced Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Shrinks the JAX-path Delaunay
pure_callbackto qhull-only and moves point location, dual areas and split points inside the JIT program. The old callback did 40-320 ms of host-serial work per likelihood evaluation (scaling with over-sampling), serialized per vmap lane (vmap_method="sequential") — the measured reason the A100 Delaunay cell runs at 106.8 ms/lane vs rectangular's 32.9 ms (autolens_profilingresults/notes/sparse_vs_dense_inversion_path.md). Expected effect: Delaunay vmap lane → ~40-50 ms (rectangular parity); to be re-profiled on the A100 when RAL returns.Point location is the same visibility-walk algorithm
find_simplexitself uses, via qhull'stri.neighbors, as a masked fixed-triplax.fori_loopchunked over query points — exact by construction. The planned 2-ring locator was implemented first and rejected by the parity gate (Δlog-evidence −0.77 on the production config; only ~31% of the relocated source-plane grid is inside the mesh hull and caustic geometry defeats ring-coverage arguments) — see the issue thread for the investigation.Known trade-off: single-thread CPU JAX-path per-eval is ~20% slower than the old callback (CPU-JAX runs the walk slower than scipy's C). GPU/vmap is the target of this refactor; the eager numpy path is byte-identical.
API Changes
None — internal changes only.
jax_delaunay/jax_delaunay_maternsignatures and return contracts (shapes, dtypes, -1 padding conventions) are unchanged; the numpy/scipy path (scipy_delaunay,scipy_delaunay_matern) is byte-identical.See full details below.
Test Plan
test_delaunay_walk.py): walk locator vsscipy.find_simplexon uniform + adaptive (blob/ring) meshes; outside-hull rows equal the KDTree nearest-vertex assignment exactly; fp edge ties must still contain the query; vertex/edge-midpoint degeneracies; adjacency tables vs the triangulationfind_simplexon the real border-relocated grids: 100.0000%autolens_workspace_testincl. a near-caustic configurationautolens_profilingruntime sweep) when RAL is backValidation checklist (--auto run)
Full API Changes (for automation & release notes)
Removed
Added (module-internal helpers, not exported via
autoarray.__init__)inversion.mesh.interpolator.delaunay.scipy_delaunay_tri_only(points_np)— qhull-only host callback returningsimplices_padded,simplex_neighbors,vertex_simplexinversion.mesh.interpolator.delaunay.pix_indexes_delaunay_walk_from(query_points, points, simplices_padded, simplex_neighbors, vertex_simplex, xp)— exact visibility-walk point locator (NumPy/JAX)DELAUNAY_LOCATE_CHUNK,DELAUNAY_WALK_STEPSChanged Behaviour
jax_delaunay/jax_delaunay_matern: identical return contract; the mappings are now computed in-JIT via the visibility walk instead of on the host. Measured identical tofind_simplexon production grids (100.0000%); the only theoretical divergence is fp ties on shared edges, which are interpolation-identical.Migration
🤖 Generated with Claude Code